# 导入必要的库和模块
from transformers import BertTokenizer, GPT2LMHeadModel, TextGenerationPipeline

# 加载模型和分词器
# 1. 使用 `GPT2LMHeadModel` 加载一个预训练的语言生成模型。
model = GPT2LMHeadModel.from_pretrained(r"D:\gpt2-chinese\models--uer--gpt2-chinese-lyric\snapshots\4a42fd76daab07d9d7ff95c816160cfb7c21684f")
# 2. 使用 `BertTokenizer` 加载与模型匹配的分词器，用于将文本转换为模型可以理解的输入格式。
tokenizer = BertTokenizer.from_pretrained(r"D:\gpt2-chinese\models--uer--gpt2-chinese-lyric\snapshots\4a42fd76daab07d9d7ff95c816160cfb7c21684f")
# 打印模型结构，方便查看模型的详细信息
print(model)

# 使用 Pipeline 调用模型
# 1. 创建一个文本生成管道 `TextGenerationPipeline`，简化了模型调用的过程。
# 2. 参数说明：
#    - model: 指定加载的模型。
#    - tokenizer: 指定加载的分词器。
#    - device="cpu": 指定运行设备为 CPU（如果支持 CUDA，可以改为 "cuda"）。
text_generator = TextGenerationPipeline(model, tokenizer, device="cpu")

# 使用 text_generator 生成文本
# 1. 循环生成 3 次歌词。
# 2. 参数说明：
#    - "当我抬起头 你正看向我,": 提供给模型的初始文本（Prompt），模型会基于这段文本继续生成后续内容。
#    - max_length=100: 设置生成文本的最大长度（包括 Prompt 部分）。
#    - do_sample=True: 是否进行随机采样。如果为 True，则每次生成的结果可能会有所不同；如果为 False，则每次都生成相同的结果。
#    - temperature=0.8: 控制生成文本的创造性。值越低，生成结果越保守；值越高，生成结果越多样化。
#    - top_p=0.9: 核采样参数，控制生成文本的质量。值越大，模型选择的词汇范围越广。
for i in range(3):
    print(text_generator("当我抬起头 你正看向我,", max_length=100, do_sample=True, temperature=0.8, top_p=0.9))